#!/usr/bin/python3
# coding=utf-8

# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import os
import shlex
import subprocess
import sys
import csv
import time
import logging
import argparse
from collections import namedtuple

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from common_scripts.gen_data import MatmulGenData
from common_scripts.compare_data import compare_data
from common_scripts.exec_utils import clear_file_cache
from common_scripts.exec_utils import get_perf_task_duration

DATA_TYPE_STR = "float16_float32"

logging.basicConfig(level=logging.INFO)

ProcessParams = namedtuple("ProcessParams", ["isProcess", "caseName", 'm', 'n', 'k', 'isBias'])


def get_case_list(work_path):
    case_list = []
    case_dir = os.path.join(work_path, "testcase")
    if not os.path.isdir(case_dir):
        logging.error(f"case path: {case_dir} not exist!")
        return case_list

    for file_name in os.listdir(case_dir):
        if not file_name.endswith(".csv"):
            continue

        abs_file_name = os.path.join(case_dir, file_name)
        with open(abs_file_name, mode='r', encoding='utf-8') as file:
            csv_reader = csv.reader(file)
            for row in csv_reader:
                item_list = [int(i) if i.strip().isdigit() else i.strip() for i in row]
                case_list.append(item_list)
    return case_list


def process_case(work_path, args, params):
    case_name, m, n, k, is_bias = (params.caseName, params.m, params.n, params.k, params.isBias)
    run_mode, is_perf, async_mode = args.r, args.p, args.m

    logging.info(f"start process case [{case_name}]")
    logging.info(f"ASYNC_MODE is set [{async_mode}], IS_PERF is set [{is_perf}]")

    clear_file_cache(work_path)

    # gen data
    MatmulGenData(m, n, k, 1, False, False, is_bias, DATA_TYPE_STR).gen_golden_data(work_path)

    params_str = f"{m} {n} {k} {is_bias}"
    if is_perf: # npu(is_perf = 1)
        cmd = f"msprof op --application=\"./ascendc_matmul_async_bbit {params_str}\" --output=./prof_out"
    elif run_mode == "sim": # sim
        cmd = f"msprof op simulator --application=\"./ascendc_matmul_async_bbit {params_str}\" --output=./sim_out"
    else: # cpu or npu(is_perf = 0)
        cmd = f"./ascendc_matmul_async_bbit {params_str}"
    subprocess.run(shlex.split(cmd))

    total_num = params.m * params.n
    if is_perf:
        wrong_num = -1
        result = None
        task_duration = get_perf_task_duration("./prof_out")
    else:
        logging.info(f"compare data case[{case_name}]")
        wrong_num = compare_data(work_path, DATA_TYPE_STR)
        result = "Fail" if (wrong_num / total_num > 0.001 or wrong_num < 0) else "Success"
        task_duration = None

    return [params.caseName, wrong_num, total_num, result, task_duration]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-r', type=str, choices=['cpu', 'sim', 'npu'], required=True)
    parser.add_argument('-p', type=int, default=0, required=False)
    parser.add_argument('-m', type=str, default='GM', choices=['GM', 'VECIN'], required=False)
    args = parser.parse_args()

    if args.r != 'npu' and args.p:
        logging.error(f"IS_PERF can only be enabled in 'npu' mode, current mode: '{args.r}'.")
        return -1

    work_path = os.path.dirname(os.getcwd())

    case_list = get_case_list(work_path)
    res_list = [["case_name", "wrong_num", "total_num", "result", "task_duration"]]

    for case in case_list:
        try:
            process_params = ProcessParams(*case)
        except Exception as e:
            logging.error(f"{case[0]}: invalid parameter, error msg: {e}")
            continue
        if process_params.isProcess:
            res_data = process_case(work_path, args, process_params)
            res_list.append(res_data)

    result_file_name = f"result_{time.time()}.csv"
    with open(os.path.join(work_path, "output", result_file_name), 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerows(res_list)

    logging.info("---------------RESULT---------------")
    for res in res_list:
        logging.info(res)
    return 0


if __name__ == "__main__":
    main()
